"""
operations for reverse image search
"""
import sys
import numpy as np
from .utils.mysql_helper import MySQL
from .utils.milvus_helper import Milvus
from .models import get_feature_extractor
from .utils.logs import LOGGER


class SearchEngine(object):
    def __init__(self, config):
        self.mysql = MySQL(config)
        self.milvus = Milvus(config)
        self.feature_extractor = get_feature_extractor(config['MODEL'])
    
    def upload(self, table_name, image_path, timestamp):
        if not self.milvus.has_collection(table_name):
            self.milvus.create_collection(table_name)
            self.milvus.create_index(table_name)
        
        try:
            feat = self.feature_extractor.extract(image_path)
            ids = self.milvus.insert(table_name, [feat], timestamp)
            
            self.mysql.create_mysql_table(table_name)
            self.mysql.load_data_to_mysql(table_name, [(str(ids[0]), image_path.encode())])
            return ids[0]
        except Exception as e:
            LOGGER.error('Error with upload: {}'.format(str(e)))
            sys.exit(1)

    def search(self, table_name, image, top_k, start_time=None, end_time=None):
        """
        search through database
        return image_paths and list of d stances.
        """
        try:
            if not self.milvus.has_collection(table_name):
                return [], [], []

            feat = self.feature_extractor.extract(image)
            vectors = self.milvus.search_vectors(table_name, [feat], top_k, start_time, end_time)
            vids = [str(x.id) for x in vectors[0]]
            if vids:
                paths = self.mysql.search_by_milvus_ids(vids, table_name)
                distances = [np.sqrt(x.distance) for x in vectors[0]]
                return paths, distances, vids
            else:
                return [], [], []
        except Exception as e:
            LOGGER.error('Error with upload: {}'.format(str(e)))
            sys.exit(1)
    
    def clear(self, table_name):
        if not self.milvus.has_collection(table_name):
            return
        self.milvus.delete_collection(table_name)
        self.mysql.delete_table(table_name)
    
    def get_database_size(self, table_name):
        n1 = self.milvus.count(table_name)
        n2 = self.mysql.count_table(table_name)
        # print('n1: {}, n2: {}'.format(n1, n2))
        return min(n1, n2)
